﻿using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;
using ParserGenerator.Extensions;

namespace ParserGenerator.DataStructures
{
    [Serializable]
    public class BinaryRelation<T> : IEnumerable<KeyValuePair<T, T>>
    {
        private readonly Dictionary<T, HashSet<T>> _dict;

        public BinaryRelation()
            : this(default(IEqualityComparer<T>))
        {
        }

        public BinaryRelation(IEqualityComparer<T> comparer)
        {
            _dict = new Dictionary<T, HashSet<T>>(comparer ?? EqualityComparer<T>.Default);
        }

        public BinaryRelation(IEnumerable<Tuple<T, T>> collection)
            : this(collection, default(IEqualityComparer<T>))
        {
        }

        public BinaryRelation(IEnumerable<Tuple<T, T>> collection, IEqualityComparer<T> comparer)
            : this(comparer)
        {
            if (collection == null)
                throw new ArgumentNullException(nameof(collection));
            AddRange(collection);
        }

        public BinaryRelation(IEnumerable<KeyValuePair<T, T>> collection)
            : this(collection, default(IEqualityComparer<T>))
        {
        }

        public BinaryRelation(IEnumerable<KeyValuePair<T, T>> collection, IEqualityComparer<T> comparer)
            : this(comparer)
        {
            if (collection == null)
                throw new ArgumentNullException(nameof(collection));
            AddRange(collection);
        }

        public BinaryRelation(BinaryRelation<T> other)
        {
            if (other == null)
                throw new ArgumentNullException(nameof(other));
            _dict = new Dictionary<T, HashSet<T>>(other._dict.Comparer);
            foreach (var t in other._dict)
                _dict.Add(t.Key, new HashSet<T>(t.Value, _dict.Comparer));
        }

        public BinaryRelation<T> Copy()
        {
            return new BinaryRelation<T>(this);
        }

        public BinaryRelation<T> Transpose()
        {
            BinaryRelation<T> result = new BinaryRelation<T>(_dict.Comparer);
            foreach (var t in _dict)
                foreach (T item in t.Value)
                    result.Add(item, t.Key);
            return result;
        }

        public void Add(T source, T target)
        {
            _dict.AddValueItem(source, target, _dict.Comparer);
        }

        public void AddRange(IEnumerable<KeyValuePair<T, T>> collection)
        {
            if (collection == null)
                throw new ArgumentNullException(nameof(collection));
            foreach (KeyValuePair<T, T> item in collection)
                Add(item.Key, item.Value);
        }

        public void AddRange(IEnumerable<Tuple<T, T>> collection)
        {
            if (collection == null)
                throw new ArgumentNullException(nameof(collection));
            foreach (Tuple<T, T> item in collection)
                Add(item.Item1, item.Item2);
        }

        public void Clear()
        {
            _dict.Clear();
        }

        private void Traverse(T source, Stack<T> stack, Dictionary<T, int> depths, Action<T, T> unionWith, Action<T, T> copyTo)
        {
            stack.Push(source);
            int depth = stack.Count;
            depths[source] = depth;
            HashSet<T> set;
            if (_dict.TryGetValue(source, out set))
            {
                foreach (T target in set)
                {
                    if (!depths.ContainsKey(target))
                        Traverse(target, stack, depths, unionWith, copyTo);
                    if (depths[target] != -1 && depths[source] > depths[target])
                        depths[source] = depths[target];
                    if (!_dict.Comparer.Equals(source, target))
                        unionWith(source, target);
                }
            }
            if (depths[source] == depth)
            {
                while (true)
                {
                    T target = stack.Pop();
                    depths[target] = -1;
                    if (_dict.Comparer.Equals(target, source))
                        break;
                    copyTo(source, target);
                }
            }
        }

        public void TransitiveClosure(Action<T, T> unionWith, Action<T, T> copyTo)
        {
            if (unionWith == null)
                throw new ArgumentNullException(nameof(unionWith));
            if (copyTo == null)
                throw new ArgumentNullException(nameof(copyTo));
            Stack<T> stack = new Stack<T>();
            Dictionary<T, int> depths = new Dictionary<T, int>(_dict.Comparer);
            foreach (var t in _dict)
                if (!depths.ContainsKey(t.Key))
                    Traverse(t.Key, stack, depths, unionWith, copyTo);
        }

        public void TransitiveClosure<TItem>(Dictionary<T, HashSet<TItem>> sets)
        {
            if (sets == null)
                throw new ArgumentNullException(nameof(sets));
            TransitiveClosure(sets, null);
        }

        public void TransitiveClosure<TItem>(Dictionary<T, HashSet<TItem>> sets, IEqualityComparer<TItem> comparer)
        {
            if (sets == null)
                throw new ArgumentNullException(nameof(sets));
            Action<T, T> unionWith = (source, target) =>
            {
                HashSet<TItem> targetSet;
                if (sets.TryGetValue(target, out targetSet))
                    sets.AddValueItems(source, targetSet, comparer);
            };
            Action<T, T> copyTo = (source, target) =>
            {
                HashSet<TItem> sourceSet;
                if (!sets.TryGetValue(source, out sourceSet))
                    sets.Remove(target);
                else
                {
                    HashSet<TItem> targetSet = sets.TryGetValueOrCreate(target, comparer);
                    targetSet.Clear();
                    targetSet.UnionWith(sourceSet);
                }
            };
            TransitiveClosure(unionWith, copyTo);
        }

        public void TransitiveClosure<TItem>(Dictionary<T, SortedSet<TItem>> sets)
        {
            if (sets == null)
                throw new ArgumentNullException(nameof(sets));
            TransitiveClosure(sets, null);
        }

        public void TransitiveClosure<TItem>(Dictionary<T, SortedSet<TItem>> sets, IComparer<TItem> comparer)
        {
            if (sets == null)
                throw new ArgumentNullException(nameof(sets));
            Action<T, T> unionWith = (source, target) =>
            {
                SortedSet<TItem> targetSet;
                if (sets.TryGetValue(target, out targetSet))
                    sets.AddValueItems(source, targetSet, comparer);
            };
            Action<T, T> copyTo = (source, target) =>
            {
                SortedSet<TItem> sourceSet;
                if (!sets.TryGetValue(source, out sourceSet))
                    sets.Remove(target);
                else
                {
                    SortedSet<TItem> targetSet = sets.TryGetValueOrCreate(target, comparer);
                    targetSet.Clear();
                    targetSet.UnionWith(sourceSet);
                }
            };
            TransitiveClosure(unionWith, copyTo);
        }

        public IEnumerator<KeyValuePair<T, T>> GetEnumerator()
        {
            foreach (var t in _dict)
                foreach (T item in t.Value)
                    yield return new KeyValuePair<T, T>(t.Key, item);
        }

        IEnumerator IEnumerable.GetEnumerator()
        {
            return GetEnumerator();
        }

        public override string ToString()
        {
            return string.Join("\n", from t in _dict select t.Key + ": {" + string.Join(", ", t.Value) + "}");
        }
    }
}
